import torch.utils.data as data
import torchvision.transforms as transforms
from PIL import Image
import numpy as np
import os

# 计算数据集的均值和标准差
def mean_std():
    c_mean = np.zeros(3)
    c_std = np.zeros(3)
    count = 0

    with open('./train.txt', mode='r', encoding='gbk') as f1:  # 使用'utf-8'编码
        for a in f1.readlines():
            img_path = a.strip().split('\t')[0]
            # 检查路径是否包含__MACOSX目录
            if '__MACOSX' in img_path:
                continue
            img = Image.open(img_path).convert('RGB')  # 确保输入的图片都是三通道
            img = np.asarray(img) / 255
            count += 1
            for d in range(3):  # 计算每一个通道的均值和标准差
                c_mean[d] += img[:, :, d].mean()
                c_std[d] += img[:, :, d].std()
    # 避免除以0
    if count == 0:
        return None, None
    mean = c_mean / count
    std = c_std / count
    return mean, std

# 定义图像转换关系
def transform(mean, std):
    transform_tran = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(mean, std)
    ])
    return transform_tran

class My_Dataset(data.Dataset):
    def __init__(self, mode, mean, std):
        super(My_Dataset, self).__init__()
        self.images = []
        self.targets = []
        self.transform = transform(mean, std)
        self.mode = mode

        with open(f'./{self.mode}.txt', mode='r', encoding='gbk') as f1:  # 使用'utf-8'编码
            for a in f1.readlines():
                img_path, target = a.strip().split('\t')
                # 检查路径是否包含__MACOSX目录
                if '__MACOSX' in img_path:
                    continue
                self.images.append(img_path)
                self.targets.append(int(target))

    def __getitem__(self, item):
        img = Image.open(self.images[item]).convert('RGB')  # 确保输入的图片都是三通道
        target = self.targets[item]
        return self.transform(img), np.asarray(target)  # 此处注意只能这样写

    def __len__(self):
        return len(self.images)

mean, std = mean_std()
if mean is not None and std is not None:
    train_dataset = My_Dataset('train', mean, std)